import os
import random
from typing import List

import torch

from data.pix2pix_dataset import Pix2pixDataset
from option import Options
from utils.util import toNormalPic


class AnimepairDataset(Pix2pixDataset):
    def hardsample(self):
        return random.random() < self.hard_reference_probability

    def get_paths(self):
        contour_dir = os.path.join(self.opt.dataroot, 'contour')
        img_dir = os.path.join(self.opt.dataroot, 'img')
        image_paths = [os.path.join(img_dir, '{:06d}.png'.format(i)) for i in self.idxs]
        contour_paths = [os.path.join(contour_dir, '{:06d}.png'.format(i)) for i in self.idxs]
        return contour_paths, image_paths

    def get_label_pt(self):
        labels = torch.load(os.path.join(self.opt.dataroot, 'labels.pt'))
        return labels

    def postprocess(self, input_dict):
        return input_dict

    def __init__(self, opt):
        super(AnimepairDataset, self).__init__(opt)


if __name__ == '__main__':
    train_parser = Options()
    opt = train_parser.parse()
    dataset = AnimepairDataset(opt)
    dict = dataset.__getitem__(0)
    sketch = dict['label']
    img = dict['image']
    ref = dict['ref']
    sketch_ref = dict['label_ref']
    similar_ref = dict['similar_ref']

    sketch = toNormalPic(sketch)
    img = toNormalPic(img)
    ref = toNormalPic(ref)
    sketch_ref = toNormalPic(sketch_ref)
    ref_similar = toNormalPic(similar_ref)

    import matplotlib.pyplot as plt
    plt.subplot(2,3,1)
    plt.imshow(sketch, cmap='gray')
    plt.subplot(2,3,2)
    plt.imshow(img)
    plt.subplot(2,3,3)
    plt.imshow(ref)
    plt.subplot(2,3,4)
    plt.imshow(sketch_ref, cmap='gray')
    plt.subplot(2,3,5)
    plt.imshow(ref_similar)
    plt.show()


